//
// Created by wkm on 2024-03-17.
//
#include <iostream>
#include "Aes.h"
#include "AddRoundKey.h"
#include "SubBytes.h"
#include "ShiftRows.h"
#include "MixColumns.h"

void helpRead(unsigned char *blk, const char *msg, int round) {
    printf("%d %s", round, msg);
    //输出密文
    for (int i = 0; i < 16; i++) {
        printf("%02x ", blk[i]);
    }
    printf("\n");
}

//AES加密函数
void AesEncrypt(unsigned char *blk, unsigned char *expansionkey, int Nr) {//加密一个区块
    //输入blk原文，直接在上面修改，输出blk密文
    //输入skey：
    //输入Nr = 10轮
    int round;
    //第1轮之前：轮密钥加
    AddRoundKey(blk, expansionkey, 0);

    helpRead(blk, "ciphertext xor:", 0);

    //第1-9轮：4类操作：字节代换、行移位、列混合、轮密钥加
    for (round = 1; round <= (Nr - 1); round++) {
        SubBytes(blk);        //输入16字节数组，直接在原数组上修改
        helpRead(blk, "ciphertext sub:", round);
        ShiftRows(blk);        //输入16字节数组，直接在原数组上修改
        helpRead(blk, "ciphertext shift:", round);
        MixColumns(blk);    //输入16字节数组，直接在原数组上修改
        helpRead(blk, "ciphertext mix:", round);
        AddRoundKey(blk, expansionkey, round);
        helpRead(blk, "ciphertext xor:", round);
    }
    //第10轮：不进行列混合
    SubBytes(blk);
    helpRead(blk, "ciphertext sub:", 10);
    ShiftRows(blk);
    helpRead(blk, "ciphertext shift:", 10);
    AddRoundKey(blk, expansionkey, Nr);
    helpRead(blk, "ciphertext xor:", 10);
}

//AES 解密函数
void Contrary_AesEncrypt(unsigned char *blk, unsigned char *expansionkey, int Nr) {
    int x;
    AddRoundKey(blk, expansionkey, Nr);
    Contrary_ShiftRows(blk);
    Contrary_SubBytes(blk);
    for (x = (Nr - 1); x >= 1; x--) {
        AddRoundKey(blk, expansionkey, x);
        Contrary_MixColumns(blk);
        Contrary_ShiftRows(blk);
        Contrary_SubBytes(blk);
    }
    AddRoundKey(blk, expansionkey, 0);
}

